// src/utility/distributed_lock.rs

use redis::{aio::MultiplexedConnection, Cmd, RedisResult, Value};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::time::{sleep, timeout};

pub struct DistributedLock {
    conn: MultiplexedConnection,
    lock_key: String,
    lock_value: String,
    expire_seconds: u64,
}

impl DistributedLock {
    /// 创建一个新的分布式锁实例
    pub fn new(
        conn: MultiplexedConnection,
        lock_key: String,
        lock_value: String,
        expire_seconds: u64,
    ) -> Self {
        Self {
            conn,
            lock_key,
            lock_value,
            expire_seconds,
        }
    }

    /// 尝试获取锁，立即返回结果
    pub async fn try_lock(&mut self) -> RedisResult<bool> {
        let mut cmd = Cmd::new();
        cmd.arg("SET")
            .arg(&self.lock_key)
            .arg(&self.lock_value)
            .arg("NX")
            .arg("EX")
            .arg(self.expire_seconds.to_string());

        let result: RedisResult<Value> = cmd.query_async(&mut self.conn).await;

        match result {
            Ok(Value::Okay) => Ok(true),
            Ok(Value::Nil) => Ok(false),
            Ok(_) => Ok(false),
            Err(e) => Err(e),
        }
    }

    /// 尝试获取锁，带重试机制
    pub async fn lock(&mut self, retry_times: u32, retry_interval_ms: u64) -> RedisResult<bool> {
        for _ in 0..retry_times {
            if self.try_lock().await.unwrap_or(false) {
                return Ok(true);
            }
            sleep(Duration::from_millis(retry_interval_ms)).await;
        }
        Ok(false)
    }

    /// 带超时的获取锁方法
    pub async fn lock_with_timeout(
        &mut self,
        timeout_duration: Duration,
    ) -> RedisResult<bool> {
        let result = timeout(timeout_duration, async {
            loop {
                if self.try_lock().await.unwrap_or(false) {
                    return true;
                }
                sleep(Duration::from_millis(100)).await;
            }
        }).await;

        match result {
            Ok(_) => Ok(true),
            Err(_) => Ok(false), // 超时
        }
    }

    /// 释放锁（只有锁的持有者才能释放）
    pub async fn unlock(&mut self) -> RedisResult<bool> {
        let script = r#"
            if redis.call("GET", KEYS[1]) == ARGV[1] then
                return redis.call("DEL", KEYS[1])
            else
                return 0
            end
        "#;

        let mut cmd = Cmd::new();
        cmd.arg("EVAL")
            .arg(script)
            .arg("1")
            .arg(&self.lock_key)
            .arg(&self.lock_value);

        let result: RedisResult<Value> = cmd.query_async(&mut self.conn).await;

        match result {
            Ok(Value::Array(values)) => {
                if let Some(Value::Int(val)) = values.first() {
                    Ok(*val > 0)
                } else {
                    Ok(false)
                }
            }
            Ok(Value::Int(val)) => Ok(val > 0),
            Ok(_) => Ok(false),
            Err(e) => Err(e),
        }
    }

    /// 续约锁的过期时间
    pub async fn renew(&mut self) -> RedisResult<bool> {
        let script = r#"
            if redis.call("GET", KEYS[1]) == ARGV[1] then
                return redis.call("EXPIRE", KEYS[1], ARGV[2])
            else
                return 0
            end
        "#;

        let mut cmd = Cmd::new();
        cmd.arg("EVAL")
            .arg(script)
            .arg("1")
            .arg(&self.lock_key)
            .arg(&self.lock_value)
            .arg(self.expire_seconds.to_string());

        let result: RedisResult<Value> = cmd.query_async(&mut self.conn).await;

        match result {
            Ok(Value::Array(values)) => {
                if let Some(Value::Int(val)) = values.first() {
                    Ok(*val > 0)
                } else {
                    Ok(false)
                }
            }
            Ok(Value::Int(val)) => Ok(val > 0),
            Ok(_) => Ok(false),
            Err(e) => Err(e),
        }
    }

    /// 检查锁是否被当前实例持有
    pub async fn is_locked(&mut self) -> RedisResult<bool> {
        let mut cmd = Cmd::new();
        cmd.arg("GET").arg(&self.lock_key);

        let result: RedisResult<Value> = cmd.query_async(&mut self.conn).await;

        match result {
            Ok(Value::BulkString(data)) => {
                let value_str = String::from_utf8_lossy(&data);
                Ok(value_str == self.lock_value)
            }
            Ok(Value::Nil) => Ok(false),
            Ok(_) => Ok(false),
            Err(e) => Err(e),
        }
    }
}

/// 便捷函数：创建分布式锁
pub fn create_lock(
    conn: MultiplexedConnection,
    resource_key: &str,
    owner_id: &str,
    expire_seconds: u64,
) -> DistributedLock {
    let lock_key = format!("lock:{}", resource_key);
    let lock_value = format!("{}:{}", owner_id, get_current_timestamp());

    DistributedLock::new(conn, lock_key, lock_value, expire_seconds)
}

/// 获取当前时间戳（毫秒）
fn get_current_timestamp() -> u128 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("Time went backwards")
        .as_millis()
}

#[cfg(test)]
mod tests {
    use super::*;
    use redis::Client;
    use tokio;

    #[tokio::test]
    async fn test_distributed_lock() -> RedisResult<()> {
        let client = Client::open("redis://127.0.0.1/")?;
        let mut conn1 = client.get_multiplexed_async_connection().await?;
        let mut conn2 = client.get_multiplexed_async_connection().await?;

        let mut lock1 = DistributedLock::new(
            conn1,
            "test_lock".to_string(),
            "owner1".to_string(),
            10,
        );

        let mut lock2 = DistributedLock::new(
            conn2,
            "test_lock".to_string(),
            "owner2".to_string(),
            10,
        );

        // 第一个客户端获取锁
        assert!(lock1.try_lock().await?);

        // 第二个客户端尝试获取同一个锁应该失败
        assert!(!lock2.try_lock().await?);

        // 释放锁
        assert!(lock1.unlock().await?);

        // 现在第二个客户端应该能获取锁
        assert!(lock2.try_lock().await?);

        // 释放锁
        assert!(lock2.unlock().await?);

        Ok(())
    }
}
